import tvm
from tvm import relax
from tvm.relax.frontend import nn
import time
from tvm.ir.instrument import PassTimingInstrument

def build_mlp_model():
    bb = relax.BlockBuilder()
    n = tvm.tir.Var("n", "int64")
    input_size = 784
    hidden_sizes = [128, 32]
    output_size = 10

    with bb.function("main"):
        model = nn.Sequential(
            nn.Linear(input_size, hidden_sizes[0]),
            nn.ReLU(),
            nn.Linear(hidden_sizes[0], hidden_sizes[1]),
            nn.ReLU(),
            nn.Linear(hidden_sizes[1], output_size),
            nn.LogSoftmax(),
        )
        # 改成使用 nn.Placeholder 声明输入，支持动态batch size n
        data = nn.Placeholder((n, input_size), name="data")
        output = model(data)
        params = [data] + model.parameters()
        bb.emit_func_output(output, params=params)
    mod = bb.get()
    return mod


if __name__ == '__main__':
    mod = build_mlp_model()
    target = "llvm"
    timing_inst = PassTimingInstrument()

    before_time = time.time()
    try:
        with tvm.transform.PassContext(opt_level=3, instruments=[timing_inst]):
            ex = relax.build(mod, target=target)
        compile_time = time.time() - before_time
        print(f"编译时间: {compile_time:.2f}秒")
        print("Pass 执行时间分析:")
        print(timing_inst.render())
    except Exception as e:
        print(f"编译失败: {str(e)}")
